#include <stdexcept>
#include "rpc_root.h"

#if defined(_MSC_VER)
#include <Windows.h>
#pragma comment(lib, "wsock32.lib")
#endif

// TODO byte order

namespace rpc {

void Rpc::initialize(Logger* logger, bool coroMode) {
    logger_ = logger;
    coroMode_ = coroMode;
}

void Rpc::update() {
    coro_sched();
    ProxyCallManagerRef.update();
    StubManagerRef.update();
}

bool Rpc::readHeader(TransportPtr transport, RpcMsgHeader& header) {
    if (transport->size() <= static_cast<int>(sizeof(RpcMsgHeader))) {
        return false;
    }
    transport->peek(reinterpret_cast<char*>(&header), sizeof(header));
    header.ntoh();
    if ((int)(header.length) > transport->size()) {
        return false;
    }
    switch (static_cast<RpcMsgType>(header.type)) {
        case RpcMsgType::RPC_CALL:
        case RpcMsgType::RPC_RETURN:
            break;
        default:
            return false;
    }
    return true;
}

void Rpc::doMessage(const RpcMsgHeader& header, TransportPtr transport, kratos::service::ServiceContext* ctx) {
    switch (static_cast<RpcMsgType>(header.type)) {
        case RpcMsgType::RPC_CALL:
            onCall(header, transport, ctx);
            break;
        case RpcMsgType::RPC_RETURN:
            onReturn(header, transport);
            break;
        default:
            break;
    }
}

void Rpc::onMessage(TransportPtr transport, kratos::service::ServiceContext* ctx) {
    while (true) {
        if (transport->isClose()) {
            return;
        }
        RpcMsgHeader header;
        if (!readHeader(transport, header)) {
            return;
        }
        auto beforeSize = transport->size();
        doMessage(header, transport, ctx);
        auto actualConsume = beforeSize - transport->size();
        if (actualConsume != static_cast<int>(header.length)) {
            if (actualConsume > static_cast<int>(header.length)) {
                transport->close();
            } else {
                transport->skip(header.length - actualConsume);
            }
        }
    }
}

StubPtr Rpc::getService(ServiceID serviceID, ServiceUUID uuid, kratos::service::ServiceContext* ctx) {
    if (serviceID) {
        auto stub = StubCreatorRef.getService(serviceID);
        if (!stub) {
            return StubCreatorRef.createService(uuid, ctx);
        } else {
            return stub;
        }
    } else {
        return StubCreatorRef.createService(uuid, ctx);
    }
}

void Rpc::onCall(const RpcMsgHeader& header, TransportPtr transport, kratos::service::ServiceContext* ctx) {
    RpcCallHeader callHeader;
    transport->recv(reinterpret_cast<char*>(&callHeader), sizeof(callHeader));
    callHeader.ntoh();
    StubPtr stub = getService(callHeader.callHeader.serviceID,
        callHeader.callHeader.serviceUUID, ctx);
    if (!stub) {
        notFound(transport, callHeader.callHeader.callID);
        return;
    }
    auto bytes = header.length - sizeof(callHeader);
    stub->callStub(transport, static_cast<int>(bytes), callHeader.callHeader);
}

void Rpc::onReturn(const RpcMsgHeader& header, TransportPtr transport) {
    RpcRetHeader retHeader;
    transport->recv(reinterpret_cast<char*>(&retHeader), sizeof(RpcRetHeader));
    retHeader.ntoh();
    auto call = ProxyCallManagerRef.get(retHeader.retHeader.callID);
    if (!call) {
        Logbegin << "ProxyCall not found [" << retHeader.retHeader.callID << "]" << Logend;
        return;
    }
    call->setErrorID(retHeader.retHeader.errorID);
    auto proxy = ProxyManagerRef.get(call->getProxyID());
    if (!proxy) {
        Logbegin << "Proxy not found [" << call->getProxyID() << "]" << Logend;
        ProxyCallManagerRef.destroy(retHeader.retHeader.callID);
        return;
    }
    auto bytes = header.length - sizeof(retHeader);
    auto errorID = static_cast<RpcError>(retHeader.retHeader.errorID);
    switch (errorID) {
        case RpcError::NOT_FOUND:
            Logbegin << "Remote service or method not found ["
                << call->getCallName() << "]" << Logend;
            proxy->setTarget(INVALID_SERVICE_ID);
            break;
        case RpcError::EXCEPTION:
            Logbegin << "Remote service's method thrown exception ["
                << call->getCallName() << "]" << Logend;
            proxy->setTarget(INVALID_SERVICE_ID);
            break;
        case RpcError::OK:
            if (proxy->isAttached()) {
                proxy->setTarget(retHeader.retHeader.serviceID);
            } else {
                proxy->setTarget(INVALID_SERVICE_ID);
            }
            break;
        default:
            proxy->setTarget(INVALID_SERVICE_ID);
            Logbegin << "Unknown RPC error " << retHeader.retHeader.errorID << " ["
                << call->getCallName() << "]" << Logend;
            ProxyCallManagerRef.destroy(call->getCallID());
            return;
    }
    call->doRet(static_cast<int>(bytes));
}

void Rpc::notFound(TransportPtr transport, CallID callID) {
    RpcRetHeader ret;
    ret.header.length = sizeof(ret);
    ret.header.type = static_cast<MsgTypeID>(RpcMsgType::RPC_RETURN);
    ret.retHeader.callID = callID;
    ret.retHeader.errorID = static_cast<ErrorID>(RpcError::NOT_FOUND);
    ret.hton();
    transport->send(reinterpret_cast<char*>(&ret), sizeof(ret));
}

Logger* Rpc::getLogger() {
    return logger_;
}

void Rpc::setContext(void* context) {
    context_ = context;
}

void buildHeader(RpcCallHeader& cheader, rpc::Proxy* proxy, rpc::CallID callID, rpc::MethodID methodID, std::uint32_t bodySize) {
    cheader.header.length = sizeof(RpcCallHeader) + bodySize;
    cheader.header.type = static_cast<rpc::MsgTypeID>(rpc::RpcMsgType::RPC_CALL);
    cheader.callHeader.serviceUUID = proxy->getServiceUUID();
    cheader.callHeader.serviceID = proxy->getTarget();
    cheader.callHeader.methodID = methodID;
    cheader.callHeader.callID = callID;
    cheader.hton();
}

LogWriter::LogWriter(Logger* logger)
    : logger_(logger) {
    ss_ << "[RPC]";
}

void initialize(Logger* logger, bool coroMode) {
    RpcRef.initialize(logger, coroMode);
}

bool loadClass(const std::string& uuid, const std::string& bundlePath, kratos::service::ServiceContext* ctx) {
    if (uuid.empty() || bundlePath.empty()) {
        return false;
    }
    return StubCreatorRef.loadClass(std::stoull(uuid), bundlePath, ctx);
}

void unloadClassUnsafe(const std::string& uuid) {
    auto& serviceMap = StubManagerRef.getMutableServiceMapUnsafe();
    for (auto& it : serviceMap) {
        Logbegin << "Destroy service [" << it.first << "][" << it.second->getServiceName() << "]" << Logend;
        it.second->reset(nullptr, &RpcRef);
    }
    serviceMap.clear();
    StubCreatorRef.unloadClass(std::stoull(uuid));
}

void coro(coroutine::CoroEntry coro) {
    coro_start(coro);
}

void deinitialize() {
    Singleton<ProxyManager>::destroy();
    Singleton<ProxyCallManager>::destroy();
    Singleton<ProxyCreator>::destroy();
    Singleton<StubCallManager>::destroy();
    Singleton<StubManager>::destroy();
    Singleton<StubCreator>::destroy();
    Singleton<Rpc>::destroy();
    //ObjectPoolManagerSinglton::destroy();
    coro_shutdown();
}

void update() {
    RpcRef.update();
}

void onMessage(TransportPtr transport, kratos::service::ServiceContext* ctx) {
    RpcRef.onMessage(transport, ctx);
}

}
